﻿#pragma once

#include <learn/alg/alg_coeff_matrix.h>
#include <learn/alg/alg_geom.h>
#include <learn/alg/alg_iterator.h>
#include <learn/geom/geom_bezier_curve.h>
#include <learn/geom/geom_bezier_surface.h>
#include <learn/geom/geom_bspline_curve.h>
#include <learn/geom/geom_bspline_surface.h>
#include <learn/math/math.h>

namespace xi
{

namespace alg
{

using namespace xi::math;
using namespace xi::geom;

template <std::size_t N> struct Alg_LaneRiesenfeld : Alg_Iterator<std::vector<VecT<N, double>>>
{
    using Vec = VecT<N, double>; // P

    std::size_t p{3};      // 样条次数
    std::vector<Vec> cpts; // 控制点

  protected:
    /**
     * @brief Lane-Riesenfeld 算法
     *
     * @return std::vector<Vec>
     */
    std::vector<Vec> next() override
    {
        auto &final = cpts;
        std::size_t n = final.size() - 1;
        std::vector<Vec> vertices(2 * n);
        for (int i = 0; i < 2 * n; i++)
        {
            if (i % 2 == 0)
                vertices[i] = final[i / 2];
            else
                vertices[i] = (final[(i - 1) / 2] + final[(i + 1) / 2]) / 2;
        }

        for (int k = 2; k < p + 1; k++)
            for (int i = 0; i < vertices.size() - 1; i++)
                vertices[i] = (vertices[i] + vertices[i + 1]) / 2;

        final = std::move(vertices);
        return final;
    }
};

template <std::size_t N> struct Alg_Bezier_Approximator
{
    using Vec = VecT<N, double>; // P

    /**
     * @brief Bezier 光顺曲线拟合
     *
     * @param[in] n
     * @return std::shared_ptr<Geom_BezierCurve<N>>
     */
    std::shared_ptr<Geom_BezierCurve<N>> approximate(auto n) const
    {
        auto m = data.size() - 1;

        // 填系数矩阵
        MatrixXd M(m + n, n + 1);
        for (std::size_t i = 0; i < m + 1; i++)
            for (std::size_t j = 0; j < n + 1; j++)
                M(i, j) = bernstein(n, j, params[i]);

        for (std::size_t i = m + 1; i < m + n; i++)
        {
            M(i, i - m - 1) = 1;
            M(i, i - m) = -2;
            M(i, i - m + 1) = 1;
        }

        // 填充 Q
        MatrixXd Q(N, m + n);
        Q.set_zero();
        for (int i = 0; i < m + 1; i++)
            for (int j = 0; j < N; j++)
                Q(j, i) = data[i][j];

        Matrix_QR_Solver solver(std::move(M));
        solver.perform();
        auto P = solver.solve(std::move(Q));

        // 构造 Bezier 曲线
        std::vector<Vec> cpts(n + 1);
        for (int i = 0; i < n + 1; i++)
            for (int j = 0; j < N; j++)
                cpts[i][j] = P(j, i);

        return std::make_shared<Geom_BezierCurve<N>>(std::move(cpts), n);
    }

    std::vector<double> params; // 参数
    std::vector<Vec> data;      // 因变量
};

template <std::size_t N> struct Alg_BSplineCurveEnergy_Approximator
{
    using Vec = VecT<N, double>; // P

    std::size_t degree{3};
    double wStretch{0.3 * 1e-8};
    double wBending{0.7 * 1e-8};
    std::vector<double> params;
    std::vector<Vec> data;

  public:
    /**
     * @brief 最小化能量的 B-spline 曲线拟合
     *
     * @param[in] n
     * @return std::shared_ptr<Geom_BSplineCurve<N>>
     */
    std::shared_ptr<Geom_BSplineCurve<N>> approximate(auto n)
    {
        n = n - 1;
        auto m = data.size() - 1;

        if (params.empty())
            params = Alg_Parameter<N>::chord(data);

        auto knots = Alg_Knot<N>::uniform(params, n + 1, degree);
        auto B0 = Alg_CoeffMatrix::curve_least_square(params, knots, degree);
        auto [B1, B2] = Alg_CoeffMatrix::curve_stretch_bending(knots, degree);
        auto wFitting = 1 - wStretch - wBending;

        // 计算系数矩阵
        Eigen::SparseMatrix<double> A = B0.transpose() * B0 * wFitting + B1 * wStretch + B2 * wBending;
        Eigen::MatrixXd Q(m + 1, N);
        for (std::size_t i = 0; i < m + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                Q(i, j) = data[i][j];
        Q = B0.transpose() * Q * wFitting;

        Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver;
        A.makeCompressed();
        solver.compute(A);
        auto P = solver.solve(Q);

        // 构造 B-spline 曲线
        std::vector<Vec> cpts(n + 1);
        for (std::size_t i = 0; i < n + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                cpts[i][j] = P(i, j);

        return std::make_shared<Geom_BSplineCurve<N>>(std::move(cpts), std::move(knots), degree);
    }
};

template <std::size_t N> struct Alg_BSplineSurfaceEnergy_Approximator
{
    using Vec = VecT<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    double wStretch{0.3 * 1e-6};
    double wBending{0.7 * 1e-6};
    std::vector<Vec2> params;
    std::vector<Vec> data;

  public:
    /**
     * @brief 最小化能量的 B-spline 曲面拟合
     *
     * @param[in] n1
     * @param[in] n2
     * @return std::shared_ptr<Geom_BSplineSurface<N>>
     */
    std::shared_ptr<Geom_BSplineSurface<N>> approximate(std::size_t n1, std::size_t n2)
    {
        n1 = n1 - 1;
        n2 = n2 - 1;
        auto m = data.size() - 1;

        if (params.empty())
            params = Alg_Parameter<N>::projection(data);

        auto uKnots = Alg_Knot<N>::equally(n1 + 1, uDegree);
        auto vKnots = Alg_Knot<N>::equally(n2 + 1, vDegree);
        auto B0 = Alg_CoeffMatrix::surface_least_square(params, uKnots, vKnots, uDegree, vDegree);
        auto [B1, B2] = Alg_CoeffMatrix::surface_stretch_bending(uKnots, vKnots, uDegree, vDegree);
        auto wFitting = 1 - wStretch - wBending;

        // 计算系数矩阵
        // Eigen::SparseMatrix<double> A = B0.transpose() * B0 * wFitting + B1 * wStretch + B2 * wBending;
        Eigen::MatrixXd A = B0.transpose() * B0 * wFitting + B1 * wStretch + B2 * wBending;
        Eigen::MatrixXd Q(m + 1, N);
        for (std::size_t i = 0; i < m + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                Q(i, j) = data[i][j];
        Q = B0.transpose() * Q * wFitting;

        // Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver;
        // A.makeCompressed();
        // solver.compute(A);
        // auto P = solver.solve(Q);
        Eigen::MatrixXd P = A.ldlt().solve(Q);

        // 构造 B-spline 曲面
        std::vector<Vec> cpts((n1 + 1) * (n2 + 1));
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto row = i * (n2 + 1) + j;
                for (std::size_t k = 0; k < N; k++)
                    cpts[row][k] = P(row, k);
            }
        }

        return std::make_shared<Geom_BSplineSurface<N>>(std::move(cpts), std::move(uKnots), std::move(vKnots), uDegree,
                                                        vDegree);
    }
};

template <std::size_t N> struct Alg_NURBSCurve_Approximator
{
    using Vec = VecT<N, double>; // P
    using VecR = VecT<N + 1, double>;

    std::size_t degree{3};
    std::vector<double> params;
    std::vector<Vec> data;

  public:
    /**
     * @brief 两步法求解 NURBS 曲线
     *
     * @param[in] n
     * @return std::shared_ptr<Geom_BSplineCurve<N, true>>
     */
    std::shared_ptr<Geom_BSplineCurve<N, true>> approximate(auto n)
    {
        n = n - 1;
        auto m = data.size() - 1;

        if (params.empty())
            params = Alg_Parameter<N>::chord(data);

        auto knots = Alg_Knot<N>::uniform(params, n + 1, degree);

        // 基函数配置矩阵
        std::vector<Eigen::Triplet<double>> tripletsB;
        for (std::size_t i = 0; i < m + 1; i++)
        {
            double t = params[i];
            auto span = find_span(knots.data(), knots.size(), degree, t);
            auto basis = basis_funcs(knots.data(), degree, span, t);

            for (std::size_t j = 0; j < degree + 1; j++)
                tripletsB.emplace_back(i, span - degree + j, basis[j]);
        }

        Eigen::SparseMatrix<double> B(m + 1, n + 1);
        B.setFromTriplets(tripletsB.begin(), tripletsB.end());

        auto X = coeff_matrices(data, B);
        std::vector<VecR> cpts(n + 1);
        for (std::size_t i = 0; i < n + 1; i++)
        {
            // 避开权重为 0 的控制点
            auto weight = X[N](i, 0);
            if (std::abs(weight) < std::numeric_limits<double>::epsilon())
                weight = 1;

            for (std::size_t j = 0; j < N; j++)
                cpts[i][j] = X[j](i, 0) / weight;

            cpts[i][N] = X[N](i, 0);
        }

        return std::make_shared<Geom_BSplineCurve<N, true>>(std::move(cpts), std::move(knots), degree);
    }

    /**
     * @brief 通过数据和配置矩阵求解目标矩阵
     *
     * @param[in] data
     * @param[in] B
     * @return std::array<Eigen::MatrixXd, N + 1>
     */
    static std::array<Eigen::MatrixXd, N + 1> coeff_matrices(const std::vector<Vec> &data,
                                                             const Eigen::SparseMatrix<double> &B)
    {
        auto m = data.size() - 1;
        auto n = B.cols() - 1;

        Eigen::SparseMatrix<double> BT = B.transpose();
        Eigen::MatrixXd Binv = Eigen::MatrixXd(BT * B).inverse();

        // 目标点的对角矩阵 X^bar, Y^bar, Z^bar
        Eigen::SparseMatrix<double> barX[N];
        for (std::size_t i = 0; i < N; i++)
            barX[i].resize(m + 1, m + 1);

        std::vector<Eigen::Triplet<double>> tripletsBarX[N];
        for (std::size_t i = 0; i < m + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                tripletsBarX[j].emplace_back(i, i, data[i][j]);

        for (std::size_t i = 0; i < N; i++)
            barX[i].setFromTriplets(tripletsBarX[i].begin(), tripletsBarX[i].end());

        // 构造权重矩阵 M = Mx + My + Mz
        // Mx = BT * bX^2 * B - (BT * bX * B) (BT * B)^-1 (BT * bX * B)
        // My = BT * bY^2 * B - (BT * bY * B) (BT * B)^-1 (BT * bY * B)
        // Mz = BT * bZ^2 * B - (BT * bZ * B) (BT * B)^-1 (BT * bZ * B)
        Eigen::SparseMatrix<double> BTbarXB[N];
        for (std::size_t i = 0; i < N; i++)
            BTbarXB[i] = BT * barX[i] * B;

        Eigen::MatrixXd M(n + 1, n + 1);
        M.setZero();
        for (std::size_t i = 0; i < N; i++)
            M += BT * barX[i] * barX[i] * B - BTbarXB[i] * Binv * BTbarXB[i];

        // 寻找 M 的最小特征值的特征向量
        // ||Mw||_2^2 = w^T * M * w / w^T * w
        // 当 w 是最小特征值对应特征向量的任意线性组合时，||Mw||_2^2 最小
        // 因此只需要找到一个特征向量即可
        Eigen::EigenSolver<Eigen::MatrixXd> es(M);

        Eigen::VectorXcd cval = es.eigenvalues();
        std::size_t min = 0;
        for (std::size_t i = 1; i < cval.size(); i++)
            if (std::abs(cval[i]) < std::abs(cval[min]))
                min = i;

        Eigen::MatrixXcd cV = es.eigenvectors();
        std::array<Eigen::MatrixXd, N + 1> X;
        X[N].resize(n + 1, 1);
        for (std::size_t i = 0; i < n + 1; i++)
            X[N](i, 0) = std::abs(cV(i, min).real());
        // 可以通过 min ||w-w_1||_2^2 的约束优化问题，寻找特征子空间中使 w 全正的解，这里直接取绝对值有点问题

        // (BT * B)X = (BT * bX * B) * w
        // (BT * B)Y = (BT * bY * B) * w
        // (BT * B)Z = (BT * bZ * B) * w
        for (std::size_t i = 0; i < N; i++)
            X[i] = Binv * BTbarXB[i] * X[N];

        return X;
    }
};

template <std::size_t N> struct Alg_NURBSSurface_Approximator
{
    using Vec = VecT<N, double>; // P
    using VecR = VecT<N + 1, double>;

    std::size_t uDegree{3}, vDegree{3};
    std::vector<Vec2> params;
    std::vector<Vec> data;

  public:
    /**
     * @brief 两步法求解 NURBS 曲面
     *
     * @param[in] n
     * @return std::shared_ptr<Geom_BSplineSurface<N, true>>
     */
    std::shared_ptr<Geom_BSplineSurface<N, true>> approximate(std::size_t n1, std::size_t n2)
    {
        n1 = n1 - 1;
        n2 = n2 - 1;
        auto m = data.size() - 1;
        auto n = (n1 + 1) * (n2 + 1) - 1;

        if (params.empty())
            params = Alg_Parameter<N>::projection(data);

        auto uKnots = Alg_Knot<N>::equally(n1 + 1, uDegree);
        auto vKnots = Alg_Knot<N>::equally(n2 + 1, vDegree);

        // 基函数配置矩阵
        std::vector<Eigen::Triplet<double>> tripletsB;
        for (std::size_t i = 0; i < m + 1; i++)
        {
            Vec2 P = params[i];
            auto uSpan = find_span(uKnots.data(), uKnots.size(), uDegree, P.x());
            auto vSpan = find_span(vKnots.data(), vKnots.size(), vDegree, P.y());

            auto ubasis = basis_funcs(uKnots.data(), uDegree, uSpan, P.x());
            auto vbasis = basis_funcs(vKnots.data(), vDegree, vSpan, P.y());

            for (std::size_t j = 0; j < uDegree + 1; j++)
                for (std::size_t k = 0; k < vDegree + 1; k++)
                    tripletsB.emplace_back(i, (uSpan - uDegree + j) * (n2 + 1) + (vSpan - vDegree + k),
                                           ubasis[j] * vbasis[k]);
        }

        Eigen::SparseMatrix<double> B(m + 1, n + 1);
        B.setFromTriplets(tripletsB.begin(), tripletsB.end());

        auto X = Alg_NURBSCurve_Approximator<N>::coeff_matrices(data, B);
        std::vector<VecR> cpts(n + 1);
        for (std::size_t i = 0; i < n + 1; i++)
        {
            // 注意避开权重为 0 的控制点
            auto weight = X[N](i, 0);
            for (std::size_t j = 0; j < N; j++)
                cpts[i][j] = X[j](i, 0) / weight;

            cpts[i][N] = X[N](i, 0);
        }

        return std::make_shared<Geom_BSplineSurface<N, true>>(std::move(cpts), std::move(uKnots), std::move(vKnots),
                                                              uDegree, vDegree);
    }
};

} // namespace alg

} // namespace xi